# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
import logging
import math
from collections import deque
from typing import List, Optional

import requests
from recordclass import recordclass

from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway import (BaseOrderGatewayInfoMutator,
                                              Order,
                                              OrderGatewayBase,
                                              OrderGatewayStatus,
                                              OrderSide,
                                              OrderType)
from coin.exchange.base.order_gateway_logger import gen_og_log_request
from coin.exchange.bitmex.kr_rest.currency import BitmexCurrency
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.bitmex.order_gateway_log import (BitmexOrderGatewayLogProvider)
from coin.exchange.bitmex.rest.api import BitmexRestApi
from coin.exchange.bitmex.ws.private_subscriber import BitmexPrivateSubscriber
from coin.exchange.util.rate_limit import RateLimiter
from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.proto.coin_query_pb2 import (AccountBalance,
                                       AccountPosition,
                                       CurrencyBalance,
                                       ProductPosition)
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig

_BITMEX_ORDER_ID_PREFIX = '-bitmex-'
_BitmexOrder = recordclass('_BitmexOrder',
                           [
                               'order',
                               'cancel_sent',
                               'amend_sent',
                               'last_amend_time',
                               'commission',
                               'trigger_price',
                               'triggered',
                               'timestamp'
                           ])
_BitmexAmend = recordclass('_BitmexAmend', ['order_id', 'price', 'leaves_qty', 'trigger_price'])


class RequestRateTracker(object):
  def __init__(self, window_size):
    self._window_size = window_size
    self._queue = deque()
    self._sum = 0

  def update_time(self, timestamp=None):
    if not timestamp:
      timestamp = get_timestamp()
    while len(self._queue) > 0:
      if timestamp - self._queue[0][0] > self._window_size:
        elem = self._queue.popleft()
        self._sum -= elem[1]
      else:
        break

  def push(self, num_of_requests=1):
    self._queue.append((get_timestamp(), num_of_requests))
    self._sum += num_of_requests
    self.update_time()

  def sum_of_requests(self):
    self.update_time()
    return self._sum


class BitmexGatewayInfo(BaseOrderGatewayInfoMutator):
  CurrencyType = BitmexCurrency
  ProductType = BitmexFuturesProduct
  BalanceType = float
  PositionType = float

  def __init__(self, currencies, products):
    super().__init__(to_list(currencies), to_list(products))
    self._logger = logging.getLogger(__name__)

    self.set_zero_balances()
    self.set_zero_positions()

    # TODO(jaewon): Make it as OKEX futures OG
    self._avg_entry_price = {product: 0. for product in self.products}

    self._margin_balance = {}

  def _update_margin_balance(self, currency: str, balance: float):
    self._margin_balance[currency] = balance

  # Margin balance includes unrealized pnl, but this is provided by exchange
  # and could be out of date or slow.
  # DO NOT USE THIS IN TRADING, ONLY ACCEPTABLE IN PNL DISPLAY.
  def get_margin_balance(self, currency: str):
    return self._margin_balance[currency]

  def get_average_entry_price(self, product):
    return self._avg_entry_price.get(product, 0.0)

  def update_bitmex_position(self, action, data):
    product = BitmexFuturesProduct.FromStrNativeProduct(data['symbol'])
    if action == 'delete':
      self.set_position(product, 0.)
      return
    position = float(data['currentQty'])
    if position == 0:
      avg_entry_price = 0.
    else:
      avg_entry_price = float(data['avgEntryPrice'])

    self.set_position(product, position)
    self._avg_entry_price[product] = avg_entry_price

  def update_bitmex_balance(self, action, data):
    assert data['currency'] == 'XBt'
    currency = BitmexCurrency.FromStrCurrency('BTC')
    if action == 'delete':
      self._logger.info('Balance deleted. %s' % data)
      self.set_balance(currency, 0.)
      self._update_margin_balance(currency.currency, 0.)
      return

    # TODO(jaewon): marginBalance vs walletBalance?
    balance = float(data['walletBalance']) / 10.**8
    self.set_balance(currency, balance)

    margin_balance = float(data['marginBalance']) / 10.**8
    self._update_margin_balance(currency.currency, margin_balance)

  def balance_as_proto(self):
    currency_balance_list = []
    for currency, balance in self.get_balances().items():
      currency_balance = CurrencyBalance(currency=currency, total=balance)
      currency_balance_list.append(currency_balance)
    return AccountBalance(exchange='Bitmex',
                          market_type="Futures",
                          each_balance=currency_balance_list)

  def position_as_proto(self):
    position_list = []
    for symbol, position in self.get_positions().items():
      product_position = ProductPosition(symbol=symbol, net_position=position)
      position_list.append(product_position)
    return AccountPosition(exchange='Bitmex', market_type="Futures", each_position=position_list)


class BitmexOrderGateway(OrderGatewayBase):
  _exchange = 'Bitmex'
  _market_type = 'Futures'
  _api_version = None

  OGInfoType = BitmexGatewayInfo

  def __init__(self,
               currencies: List[BitmexCurrency],
               products: List[BitmexFuturesProduct],
               config,
               og_config=None,
               logger=None):
    super().__init__()
    self._logger = logger or logging.getLogger(__name__)

    default_og_config = OrderGatewayConfig(order_rate_limit=300, order_rate_limit_per=300)
    if og_config is not None:
      default_og_config.MergeFrom(og_config)
    self._og_config = default_og_config

    self._ioloop = None
    self._api = BitmexRestApi(config.api_key, config.api_secret, logger=self._logger)
    self._private_sub = BitmexPrivateSubscriber(config.api_key,
                                                config.api_secret,
                                                logger=self._logger)
    self._private_sub.on_ready_callback = self._on_private_subscriber_ready
    self._private_sub.callbacks.append(self._on_table_update)
    self._private_sub_ready = False

    self._rate_tracker = RequestRateTracker(300 * 10**9)  # 5 minutes
    self._order_rate_limiter = RateLimiter(self._og_config.order_rate_limit,
                                           self._og_config.order_rate_limit_per)
    self._orders = {}
    self._og_info = BitmexGatewayInfo(currencies, products)
    self._og_logger = BitmexOrderGatewayLogProvider(gen_og_log_request(config, self))

  def submittable(self):
    return self.is_ready() and self._order_rate_limiter.check()

  def get_average_entry_price(self, product):
    return self._og_info.get_average_entry_price(product)

  def _on_private_subscriber_ready(self, ready):
    if ready:
      self._private_sub_ready = True
    else:
      self._private_sub_ready = False
      self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)

  def _on_order(self, action, table, key, data):
    if action != 'partial':
      return

    if data['ordStatus'] not in ['New', 'PartiallyFilled']:
      return

    timestamp = get_timestamp()
    order_id = data['clOrdID'] or (_BITMEX_ORDER_ID_PREFIX + data['orderID'])
    symbol = data['symbol']
    product = BitmexFuturesProduct.FromStrNativeProduct(data['symbol'])
    order_side = OrderSide.BUY if data['side'] == 'Buy' else OrderSide.SELL
    order_type = (OrderType.LIMIT if data['ordType'] in ('Limit',
                                                         'StopLimit') else OrderType.UNKNOWN)
    if data['ordType'] == 'StopLimit':
      order_trigger_price = float(data['stopPx'])
    else:
      order_trigger_price = None
    order_triggered = bool(data['workingIndicator'])

    order = Order(order_id=order_id,
                  product=product,
                  price=float(data['price']),
                  qty=float(data['orderQty']),
                  order_side=order_side,
                  order_type=order_type,
                  tag=data.get('text', None),
                  accepted=True,
                  cancel_sent=False,
                  internal=_BitmexOrder(order=None,
                                        cancel_sent=False,
                                        amend_sent=False,
                                        last_amend_time=0,
                                        commission=None,
                                        trigger_price=order_trigger_price,
                                        triggered=order_triggered,
                                        timestamp=timestamp))
    self._orders[order_id] = order

    self._og_logger.gen_order_event_and_log(data, get_timestamp(), OrderEvent.ORDER_ACCEPTED)

  def _on_execution(self, action, table, key, data):
    if action == 'delete':
      return

    order_id = data['clOrdID'] or (_BITMEX_ORDER_ID_PREFIX + data['orderID'])
    symbol = data['symbol']
    product = BitmexFuturesProduct.FromStrNativeProduct(data['symbol'])
    timestamp = get_timestamp()

    if order_id in self._orders:
      order = self._orders[order_id]
    else:
      order_side = OrderSide.BUY if data['side'] == 'Buy' else OrderSide.SELL
      order_type = (OrderType.LIMIT if data['ordType'] in ('Limit',
                                                           'StopLimit') else OrderType.UNKNOWN)
      if data['ordType'] == 'StopLimit':
        order_trigger_price = float(data['stopPx'])
      else:
        order_trigger_price = None
      order_triggered = bool(data['workingIndicator'])

      order = Order(order_id=order_id,
                    product=product,
                    price=float(data['price']),
                    qty=float(data['orderQty']),
                    order_side=order_side,
                    order_type=order_type,
                    tag=data.get('text', None),
                    accepted=True,
                    cancel_sent=False,
                    internal=_BitmexOrder(order=None,
                                          cancel_sent=False,
                                          amend_sent=False,
                                          last_amend_time=0,
                                          commission=None,
                                          trigger_price=order_trigger_price,
                                          triggered=order_triggered,
                                          timestamp=timestamp))
      self._orders[order_id] = order

    if not order.accepted:
      order.accepted = True
      self._og_logger.gen_order_event_and_log(data, timestamp, OrderEvent.ORDER_ACCEPTED)

    stop_limit_order = (data['ordType'] == 'StopLimit')
    order_triggered = bool(data['workingIndicator'])
    if not order.internal.triggered and order_triggered:
      order.internal.triggered = order_triggered
      if stop_limit_order:
        self._logger.info('TRIGGER %s %4s %6f @ %.1f tag="%s"',
                          symbol,
                          str(order.order_side)[10:],
                          order.qty,
                          order.price,
                          order.tag or '')

    exec_type = data['execType']
    if exec_type in ['New', 'Canceled', 'Rejected']:
      pass

    elif exec_type == 'Trade':
      # Trade
      fill_price = float(data['lastPx'])
      fill_qty = float(data['lastQty'])
      self._og_logger.gen_order_event_and_log(data, timestamp, OrderEvent.ORDER_FILLED)
      if 'commission' in data:
        order.internal.commission = float(data['commission'])

      self._publish_fill(fill_price, fill_qty, order)
      self._logger.info('FILL %s %4s %6f @ %.1f tag="%s"',
                        symbol,
                        str(order.order_side)[10:],
                        fill_qty,
                        fill_price,
                        order.tag or '')

    elif exec_type == 'Replaced':
      # Amend
      order.price = float(data['price'])
      order.qty = float(data['orderQty'])
      if stop_limit_order:
        order.internal.trigger_price = float(data['stopPx'])

      order.internal.amend_sent = False
      self._og_logger.gen_order_event_and_log(data, timestamp, OrderEvent.ORDER_AMEND_ACCEPTED)

    elif exec_type == 'TriggeredOrActivatedBySystem':
      order.internal.triggered = True

    else:
      self._logger.warning('Unknown exec_type: %s' % exec_type)

    ord_status = data['ordStatus']
    if ord_status in ['Canceled', 'Filled', 'Rejected']:
      self._og_logger.gen_order_event_and_log(data, timestamp)
      del self._orders[order.order_id]

  def _on_position(self, action, table, key, data):
    self._og_info.update_bitmex_position(action, data)
    self._og_logger.write_position(self._og_info.position_as_proto())

  def _on_margin(self, action, table, key, data):
    self._og_info.update_bitmex_balance(action, data)
    self._og_logger.write_balance(self._og_info.balance_as_proto())
    if self._private_sub_ready and not self._og_ready.ready:
      self._logger.info('OG ready')
      self._og_ready.set_status(OrderGatewayStatus.READY)

  def _on_table_update(self, action, table, key, data):
    if table == 'execution':
      self._on_execution(action, table, key, data)
    elif table == 'order':
      self._on_order(action, table, key, data)
    elif table == 'position':
      self._on_position(action, table, key, data)
    elif table == 'margin':
      self._on_margin(action, table, key, data)

  def start(self, queue, ioloop, flow_sub=None):
    self._ioloop = ioloop
    self._og_ready.set_status(OrderGatewayStatus.INITIALIZING)
    self._private_sub.start(self._ioloop)
    # Reserve connection pool.
    for i in range(2):
      self._ioloop.add_callback(self._api.get_position)

  def make_bitmex_order(self,
                        product: BitmexFuturesProduct,
                        price: float,
                        qty: float,
                        order_side: OrderSide,
                        order_type: OrderType,
                        order_id: Optional[str] = None,
                        tag=None,
                        post_only: bool = False,
                        display_qty: Optional[float] = None,
                        trigger_price: Optional[float] = None):
    # TODO(jaewon): Support ice-berg order
    assert display_qty is None or display_qty == qty, display_qty

    timestamp = get_timestamp()
    bitmex_order_side = 'Buy' if order_side == OrderSide.BUY else 'Sell'
    if trigger_price is None:
      bitmex_order_type = 'Limit'
    else:
      bitmex_order_type = 'StopLimit'
    bitmex_exec_inst = 'ParticipateDoNotInitiate' if post_only else None
    bitmex_text = None if tag is None else str(tag)

    if not order_id:
      order_id = 'order-%d' % get_timestamp()
    assert not order_id.startswith(_BITMEX_ORDER_ID_PREFIX)

    bitmex_order_dict = self._api.make_order_dict(symbol=product.native_symbol,
                                                  side=bitmex_order_side,
                                                  price=price,
                                                  qty=qty,
                                                  order_type=bitmex_order_type,
                                                  exec_inst=bitmex_exec_inst,
                                                  client_order_id=order_id,
                                                  text=bitmex_text,
                                                  trigger_price=trigger_price)

    order = Order(order_id=order_id,
                  product=product,
                  price=price,
                  qty=qty,
                  order_side=order_side,
                  order_type=order_type,
                  tag=tag,
                  accepted=False,
                  cancel_sent=False,
                  internal=_BitmexOrder(order=bitmex_order_dict,
                                        cancel_sent=False,
                                        amend_sent=False,
                                        last_amend_time=0,
                                        commission=None,
                                        trigger_price=trigger_price,
                                        triggered=False,
                                        timestamp=timestamp))
    return order

  def get_working_order(self) -> List[Order]:
    return [order for order in self._orders.values()]

  def get_rate_limit(self):
    return self._api.get_rate_limit()

  def get_rate_limit_remaining(self):
    return self._api.get_rate_limit_remaining()

  def get_rate_count(self):
    return self._rate_tracker.sum_of_requests()

  def on_submit_response(self, orders, future):
    try:
      future.result()
      for order in to_list(orders):
        self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                                   True,
                                                                   order.order_id,
                                                                   get_timestamp())
    except requests.RequestException:
      for order in to_list(orders):
        self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                                   False,
                                                                   order.order_id,
                                                                   get_timestamp())
        self._logger.exception('Order error: %s' % str(order))
        del self._orders[order.order_id]

  def submit_impl(self,
                  product: BitmexFuturesProduct,
                  price: float,
                  qty: float,
                  order_side: OrderSide,
                  order_type: OrderType,
                  *,
                  order_id: Optional[str] = None,
                  tag=None,
                  post_only: bool = False,
                  display_qty: Optional[float] = None,
                  trigger_price: Optional[float] = None):
    order = self.make_bitmex_order(product=product,
                                   price=price,
                                   qty=qty,
                                   order_side=order_side,
                                   order_type=order_type,
                                   order_id=order_id,
                                   tag=tag,
                                   post_only=post_only,
                                   display_qty=display_qty,
                                   trigger_price=trigger_price)
    assert order.order_id not in self._orders
    timestamp = get_timestamp()
    self._orders[order.order_id] = order

    future = self._api.create_bulk_order([order.internal.order])
    self._ioloop.add_future(future, functools.partial(self.on_submit_response, [order]))

    self._og_logger.gen_order_request_and_order_event_and_log(order.product.symbol,
                                                              order.price,
                                                              order.qty,
                                                              order.order_side,
                                                              order.order_id,
                                                              timestamp)
    self._rate_tracker.push(1)

  def submit_bulk(self, orders):
    orders = to_list(orders)
    for order in orders:
      assert order.order_id not in self._orders
      self._orders[order.order_id] = order

    bitmex_orders = [order.internal.order for order in orders]
    future = self._api.create_bulk_order(bitmex_orders)
    self._ioloop.add_future(future, functools.partial(self.on_submit_response, orders))

    timestamp = get_timestamp()
    for order in orders:
      self._og_logger.gen_order_request_and_order_event_and_log(order.product.symbol,
                                                                order.price,
                                                                order.qty,
                                                                order.order_side,
                                                                order.order_id,
                                                                timestamp)
    self._rate_tracker.push(math.ceil(len(orders) / 10.0))

  def _make_bitmex_amend_dict(self, amend_obj: _BitmexAmend):
    order_id = amend_obj.order_id
    if order_id.startswith(_BITMEX_ORDER_ID_PREFIX):
      exchange_order_id = order_id[len(_BITMEX_ORDER_ID_PREFIX):]
      client_order_id = None
    else:
      exchange_order_id = None
      client_order_id = order_id
    amend_dict = self._api.make_amend_dict(order_id=exchange_order_id,
                                           client_order_id=client_order_id,
                                           price=amend_obj.price,
                                           leaves_qty=amend_obj.leaves_qty,
                                           trigger_price=amend_obj.trigger_price)
    return amend_dict

  def make_bitmex_amend(self,
                        order_id: str,
                        price: Optional[float] = None,
                        leaves_qty: Optional[float] = None,
                        trigger_price: Optional[float] = None):
    return _BitmexAmend(order_id=order_id,
                        price=price,
                        leaves_qty=leaves_qty,
                        trigger_price=trigger_price)

  def amend(self,
            order_id: str,
            price: Optional[float] = None,
            leaves_qty: Optional[float] = None,
            trigger_price: Optional[float] = None):
    amend_obj = self.make_bitmex_amend(order_id, price, leaves_qty, trigger_price)
    self.amend_bulk([amend_obj])

  def amend_bulk(self, amend_reqs):
    amend_reqs = to_list(amend_reqs)
    amend_dicts = []
    timestamp = get_timestamp()
    for amend_obj in amend_reqs:
      amend_dicts.append(self._make_bitmex_amend_dict(amend_obj))
      if amend_obj.order_id in self._orders:
        order = self._orders[amend_obj.order_id]
        order.internal.amend_sent = True
        order.internal.last_amend_time = timestamp
    self._api.amend_bulk_order(amend_dicts)
    for amend_obj in amend_reqs:
      if amend_obj.order_id in self._orders:
        order = self._orders[amend_obj.order_id]
        price = amend_obj.price
        if price is None:
          price = order.price
        self._og_logger.gen_order_request_and_order_event_and_log(order.product.symbol,
                                                                  price,
                                                                  amend_obj.leaves_qty,
                                                                  order.order_side,
                                                                  order.order_id,
                                                                  timestamp,
                                                                  True)
        self._og_logger.gen_order_response_and_order_event_and_log(order.product.symbol,
                                                                   True,
                                                                   order.order_id,
                                                                   timestamp,
                                                                   True)
    self._rate_tracker.push(math.ceil(len(amend_dicts) / 10.0))

  def on_cancel_response(self, order_ids, future):
    try:
      order_info_list = to_list(future.result())
      for order_info in order_info_list:
        if 'error' in order_info:
          self._logger.error('Cancel error %s', order_info['error'])

        if order_info['ordStatus'] == 'Canceled':
          order_id = order_info['clOrdID']
          order = self._orders.pop(order_id, None)
          if order is None:
            self._logger.warning(
                'Cancel confirmation received already (from WS) '
                'or it is from an unknown order %s',
                order_id)
            continue
          self._og_logger.gen_cancel_response_and_log(order.product.symbol, True, order.order_id)
        elif order_info['ordStatus'] == 'Filled':
          self._logger.info('Order %s is filled already. Cancellation failed',
                            order_info['clOrdID'])
        else:
          order.cancel_sent = False
          order.internal.cancel_sent = False
    except requests.RequestException:
      for order_id in order_ids:
        if order_id in self._orders:
          order = self._orders[order_id]
          self._og_logger.gen_cancel_response_and_log(order.product.symbol, False, order.order_id)
          # If failed, reset cancel_sent flag.
          order.cancel_sent = False
          order.internal.cancel_sent = False
          self._logger.exception('Cancel request exception (%s)', order)
        else:
          self._logger.warning('Failed to cancel %s that is not in WO!', order_id)

  def cancel_impl(self, order_id, *, ignore_error=False):
    if order_id in self._orders:
      order = self._orders[order_id]
      if order.cancel_sent is False:
        order.cancel_sent = True
        order.internal.cancel_sent = True
      else:
        self._logger.info('Cancel has been sent already: %s', order.order_id)
        return

    if order_id.startswith(_BITMEX_ORDER_ID_PREFIX):
      exchange_order_id = order_id[len(_BITMEX_ORDER_ID_PREFIX):]
      client_order_id = None
    else:
      exchange_order_id = None
      client_order_id = order_id

    future = self._api.cancel_order(order_id=exchange_order_id, client_order_id=client_order_id)
    self._ioloop.add_future(future, functools.partial(self.on_cancel_response, [order_id]))

    if order_id in self._orders:
      product = self._orders[order_id].product
      self._og_logger.gen_cancel_request_and_order_event_and_log(product.symbol,
                                                                 order_id,
                                                                 get_timestamp())
    self._rate_tracker.push(1)

  def cancel_all_impl(self, *, ignore_error=False):
    for order in self._orders.values():
      order.cancel_sent = True
      order.internal.cancel_sent = True
    future = self._api.cancel_all()
    self._ioloop.add_future(future, functools.partial(self.on_cancel_response, None))
    # Bulk cancels, regardless of count, always only count as one request.
    self._rate_tracker.push(1)

  def cancel_multiple_impl(self, order_ids, *, ignore_error=False):
    exchange_order_id_list = []
    client_order_id_list = []
    for order_id in order_ids:
      if order_id in self._orders:
        order = self._orders[order_id]
        if order.cancel_sent is False:
          order.cancel_sent = True
          order.internal.cancel_sent = True
        else:
          continue

      if order_id.startswith(_BITMEX_ORDER_ID_PREFIX):
        exchange_order_id = order_id[len(_BITMEX_ORDER_ID_PREFIX):]
        exchange_order_id_list.append(exchange_order_id)
      else:
        client_order_id = order_id
        client_order_id_list.append(client_order_id)

    if exchange_order_id_list:
      future = self._api.cancel_bulk_order(order_id=exchange_order_id_list, client_order_id=None)
      self._ioloop.add_future(future, functools.partial(self.on_cancel_response, None))
      self._rate_tracker.push(1)

    if client_order_id_list:
      future = self._api.cancel_bulk_order(order_id=None, client_order_id=client_order_id_list)
      self._ioloop.add_future(future, functools.partial(self.on_cancel_response, None))
      self._rate_tracker.push(1)

    for order_id in order_ids:
      if order_id in self._orders:
        product = self._orders[order_id].product
        self._og_logger.gen_cancel_request_and_order_event_and_log(product.symbol,
                                                                   order_id,
                                                                   get_timestamp())
